package com.wechat.back.utils;

import java.util.*;
import java.util.stream.Collectors;
public class TFIDFUtil {
    private Map<String, Double> idfMap;
    private Map<String, Map<String, Double>> tfMap;

    public TFIDFUtil(List<String> documents) {
        // 计算IDF
        Map<String, Integer> wordDocCount = new HashMap<>();
        for (String doc : documents) {
            Set<String> words = new HashSet<>(Arrays.asList(doc.split(" ")));
            for (String word : words) {
                wordDocCount.put(word, wordDocCount.getOrDefault(word, 0) + 1);
            }
        }

        int totalDocs = documents.size();
        idfMap = wordDocCount.entrySet().stream()
                .collect(Collectors.toMap(
                        Map.Entry::getKey,
                        entry -> Math.log((double) totalDocs / entry.getValue())
                ));

        // 计算TF
        tfMap = new HashMap<>();
        for (String doc : documents) {
            Map<String, Integer> wordCount = new HashMap<>();
            for (String word : doc.split(" ")) {
                wordCount.put(word, wordCount.getOrDefault(word, 0) + 1);
            }

            Map<String, Double> tf = wordCount.entrySet().stream()
                    .collect(Collectors.toMap(
                            Map.Entry::getKey,
                            entry -> (double) entry.getValue() / doc.split(" ").length
                    ));

            tfMap.put(doc, tf);
        }
    }

    public Map<String, Double> getTFIDFVector(String document) {
        Map<String, Double> tfidf = new HashMap<>();
        Map<String, Double> tf = tfMap.get(document);
        for (String word : tf.keySet()) {
            tfidf.put(word, tf.get(word) * idfMap.getOrDefault(word, 0.0));
        }
        return tfidf;
    }

    public double cosineSimilarity(Map<String, Double> vector1, Map<String, Double> vector2) {
        Set<String> keys = new HashSet<>(vector1.keySet());
        keys.addAll(vector2.keySet());

        double dotProduct = 0.0;
        double normA = 0.0;
        double normB = 0.0;

        for (String key : keys) {
            double val1 = vector1.getOrDefault(key, 0.0);
            double val2 = vector2.getOrDefault(key, 0.0);

            dotProduct += val1 * val2;
            normA += Math.pow(val1, 2);
            normB += Math.pow(val2, 2);
        }

        return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
    }
}